package app

import (
	"bytes"
	"fmt"
	"net"
	"smb-trans/common"
	"strconv"
	"time"
)

type SnapConn struct {
	net.Conn
	snapBuffer *bytes.Buffer
}

func (conn *SnapConn) Read(b []byte) (n int, err error) {
	if conn.snapBuffer != nil {
		n, err := conn.snapBuffer.Read(b)
		if err == nil {
			if conn.snapBuffer.Len() == 0 {
				conn.snapBuffer = nil
			}
		}
		return n, err
	}
	return conn.Conn.Read(b)

}

func FindAddrHandler(conn net.Conn, readTime int) (net.Conn, string, error) {
	//return conn, "", nil
	cacheConn := &SnapConn{
		Conn:       conn,
		snapBuffer: new(bytes.Buffer),
	}
	buffer := make([]byte, 2024)
	_ = conn.SetReadDeadline(time.Now().Add(time.Second))
	defer func() {
		_ = conn.SetReadDeadline(time.Time{})
	}()
	n, err := conn.Read(buffer)
	if err != nil {
		return conn, "", err
	}
	cacheConn.snapBuffer.Write(buffer[:n])
	res, ok := common.TryFindHost(buffer[:n])
	if ok {
		return cacheConn, res, nil
	}
	fmt.Println(string(buffer[:n]), "readTime:"+strconv.Itoa(readTime))
	return cacheConn, "", nil
}
